package com.tpp.threat_perception_platform.consumer;

import com.alibaba.dashscope.aigc.generation.GenerationResult;
import com.alibaba.fastjson.JSON;
import com.rabbitmq.client.Channel;

//import com.tpp.threat_perception_platform.pojo.Host;
//import com.tpp.threat_perception_platform.service.HostService;
import com.tpp.threat_perception_platform.pojo.*;
import com.tpp.threat_perception_platform.pojo.Process;
import com.tpp.threat_perception_platform.service.*;
import com.tpp.threat_perception_platform.utils.Ai;
import com.tpp.threat_perception_platform.utils.EncryptionUtils;
import org.springframework.amqp.rabbit.annotation.RabbitListener;
import org.springframework.amqp.support.AmqpHeaders;
//import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.beans.factory.annotation.Autowired;
import org.springframework.messaging.handler.annotation.Headers;
import org.springframework.stereotype.Component;

import java.io.IOException;
import java.util.HashSet;
import java.util.List;
import java.util.Map;
import java.util.Set;

@Component
public class RabbitSysInfoConsumer {

    @Autowired
    private HostService hostService;

    @Autowired
    private AccountService accountService;

    @Autowired
    private ServiceService serviceService;

    @Autowired
    private AppService appService;

    @Autowired
    private ProcessService processService;

    @Autowired
    private HotfixService hotfixService;

    @Autowired
    private WeakPwdService weakPwdService;

    @Autowired
    private HostAppVulService hostAppVulService;

    @Autowired
    private HostVulService hostVulService;

    @Autowired
    private HostSystemVulService hostSystemVulService;

    @Autowired
    private HostLogService hostLogService;

    @RabbitListener(queues = "sysinfo_queue")

    public void receive(String message, @Headers Map<String,Object> headers, Channel channel) throws IOException {
        System.out.println("Received message: " + message);

        // 解密数据
        try {
            message = EncryptionUtils.decrypt(message);
            System.out.println("Decrypted message: " + message);
        } catch (Exception e) {
            System.out.println("解密失败");
        }

//         将数据存储到数据库
        try {
//            System.out.println("Received message: " + message);
            Host host = JSON.parseObject(message, Host.class);
            // 存储到数据库
            int res = hostService.saveHost(host);

            if (res > 0){
                // 手动 ACK, 先获取 deliveryTag
                Long deliveryTag = (Long)headers.get(AmqpHeaders.DELIVERY_TAG);
                // ACK
                channel.basicAck(deliveryTag,false);
            }
        } catch (Exception e) {
            // throw new RuntimeException(e);
            // 手动 ACK, 先获取 deliveryTag
            Long deliveryTag = (Long)headers.get(AmqpHeaders.DELIVERY_TAG);
            // ACK
            channel.basicAck(deliveryTag,false);
        }
    }

    @RabbitListener(queues = "status_queue")
    public void receiveStatus(String message, @Headers Map<String,Object> headers, Channel channel) throws IOException {
        System.out.println("Received message: " + message);

        // 解密数据
        try {
            message = EncryptionUtils.decrypt(message);
            System.out.println("Decrypted message: " + message);
        } catch (Exception e) {
            System.out.println("解密失败");
        }

        //反序列化数据、
        try {
            Host host = JSON.parseObject(message, Host.class);
            //根据macAddress更新数据
            int res=hostService.updateHostByMacAddress(host);
            if (res > 0){
                // 手动 ACK, 先获取 deliveryTag
                Long deliveryTag = (Long)headers.get(AmqpHeaders.DELIVERY_TAG);
                // ACK
                channel.basicAck(deliveryTag,false);
            }
        } catch (Exception e) {
            Long deliveryTag = (Long)headers.get(AmqpHeaders.DELIVERY_TAG);
            // ACK
            channel.basicAck(deliveryTag,false);
        }
    }

    @RabbitListener(queues = "account_queue")
    public void receiveAccount(String message, @Headers Map<String, Object> headers, Channel channel) throws IOException {
        System.out.println("Received message: " + message);

        // 解密数据
        try {
            message = EncryptionUtils.decrypt(message);
            System.out.println("Decrypted message: " + message);
        } catch (Exception e) {
            System.out.println("解密失败");
        }

        // 获取消息的deliveryTag，用于后续确认
        Long deliveryTag = (Long) headers.get(AmqpHeaders.DELIVERY_TAG);


        try {
            List<Account> accounts = JSON.parseArray(message, Account.class);

            boolean allSuccess = true;

            // 删除数据库中已存在的账户对象

            Set<String> macAddresses = new HashSet<>();
            for (Account account : accounts) {
                macAddresses.add(account.getMacAddress());
            }
            // 批量删除这些macAddress对应的记录
            for (String macAddress : macAddresses) {
                accountService.deleteByMacAddress(macAddress);
            }


            // 处理所有账户对象
            GenerationResult result = Ai.accountAi(accounts);
            String analyzedJson = result.getOutput().getChoices().get(0).getMessage().getContent();
            List<Account> analyzedAccounts = JSON.parseArray(analyzedJson, Account.class);
            for (Account account : analyzedAccounts) {
                try {
                    // 存储到数据库
                    int res = accountService.saveAccount(account);
                    if (res <= 0) {
                        allSuccess = false;
                        System.err.println("Failed to save account: " + account);
                    }
                } catch (Exception e) {
                    allSuccess = false;

                }
            }

            // 所有账户处理完成后，统一进行ACK确认
            if (allSuccess) {
                channel.basicAck(deliveryTag, false);
            } else {
                // 处理失败，根据业务需求选择拒绝消息或重试
                channel.basicReject(deliveryTag, false); // false表示不重新入队
            }
        } catch (Exception e) {
            // 处理消息解析或其他严重错误

            channel.basicReject(deliveryTag, false); // 拒绝消息，不重新入队
        }
    }

    @RabbitListener(queues = "service_queue")
    public void receiveService(String message, @Headers Map<String, Object> headers, Channel channel) throws IOException {
        System.out.println("Received message: " + message);

        // 解密数据
        try {
            message = EncryptionUtils.decrypt(message);
            System.out.println("Decrypted message: " + message);
        } catch (Exception e) {
            System.out.println("解密失败");
        }

        // 获取消息的deliveryTag，用于后续确认
        Long deliveryTag = (Long) headers.get(AmqpHeaders.DELIVERY_TAG);


        try {
            List<Service> services = JSON.parseArray(message, Service.class);
            boolean allSuccess = true;

            // 删除数据库中已存在的账户对象

            Set<String> macAddresses = new HashSet<>();
            for (Service service : services) {
                macAddresses.add(service.getMacAddress());
            }
            // 批量删除这些macAddress对应的记录
            for (String macAddress : macAddresses) {
                serviceService.deleteByMacAddress(macAddress);
            }


            // 处理所有账户对象
            GenerationResult result = Ai.serviceAi(services);
            String analyzedJson = result.getOutput().getChoices().get(0).getMessage().getContent();
            List<Service> analyzedServices = JSON.parseArray(analyzedJson, Service.class);
            for (Service service : analyzedServices) {
                try {
                    // 存储到数据库
                    int res = serviceService.saveService(service);
                    if (res <= 0) {
                        allSuccess = false;
                        System.err.println("Failed to save service: " + service);
                    }
                } catch (Exception e) {
                    allSuccess = false;

                }
            }

            // 所有账户处理完成后，统一进行ACK确认
            if (allSuccess) {
                channel.basicAck(deliveryTag, false);
            } else {
                // 处理失败，根据业务需求选择拒绝消息或重试
                channel.basicReject(deliveryTag, false); // false表示不重新入队
            }
        } catch (Exception e) {
            // 处理消息解析或其他严重错误

            channel.basicReject(deliveryTag, false); // 拒绝消息，不重新入队
        }
    }

    @RabbitListener(queues = "app_queue")
    public void receiveApp(String message, @Headers Map<String, Object> headers, Channel channel) throws IOException {
        System.out.println("Received message: " + message);

        // 解密数据
        try {
            message = EncryptionUtils.decrypt(message);
            System.out.println("Decrypted message: " + message);
        } catch (Exception e) {
            System.out.println("解密失败");
        }

        // 获取消息的deliveryTag，用于后续确认
        Long deliveryTag = (Long) headers.get(AmqpHeaders.DELIVERY_TAG);


        try {
            List<App> apps = JSON.parseArray(message, App.class);
            boolean allSuccess = true;

            // 删除数据库中已存在的账户对象

            Set<String> macAddresses = new HashSet<>();
            for (App app : apps) {
                macAddresses.add(app.getMacAddress());
            }
            // 批量删除这些macAddress对应的记录
            for (String macAddress : macAddresses) {
                appService.deleteByMacAddress(macAddress);
            }


            // 处理所有账户对象
            for (App app : apps) {
                try {
                    // 存储到数据库
                    int res = appService.saveApp(app);
                    if (res <= 0) {
                        allSuccess = false;
                        System.err.println("Failed to save app: " + app);
                    }
                } catch (Exception e) {
                    allSuccess = false;

                }
            }

            // 所有账户处理完成后，统一进行ACK确认
            if (allSuccess) {
                channel.basicAck(deliveryTag, false);
            } else {
                // 处理失败，根据业务需求选择拒绝消息或重试
                channel.basicReject(deliveryTag, false); // false表示不重新入队
            }
        } catch (Exception e) {
            // 处理消息解析或其他严重错误

            channel.basicReject(deliveryTag, false); // 拒绝消息，不重新入队
        }
    }


    @RabbitListener(queues = "process_queue")
    public void receiveProcess(String message, @Headers Map<String, Object> headers, Channel channel) throws IOException {
        System.out.println("Received message: " + message);

        // 解密数据
        try {
            message = EncryptionUtils.decrypt(message);
            System.out.println("Decrypted message: " + message);
        } catch (Exception e) {
            System.out.println("解密失败");
        }

        // 获取消息的deliveryTag，用于后续确认
        Long deliveryTag = (Long) headers.get(AmqpHeaders.DELIVERY_TAG);


        try {
            List<Process> processes = JSON.parseArray(message, Process.class);
            boolean allSuccess = true;

            // 删除数据库中已存在的账户对象

            Set<String> macAddresses = new HashSet<>();
            for (Process process : processes) {
                macAddresses.add(process.getMacAddress());
            }
            // 批量删除这些macAddress对应的记录
            for (String macAddress : macAddresses) {
                processService.deleteByMacAddress(macAddress);
            }


            // 处理所有账户对象
            for (Process process : processes) {
                try {
                    // 存储到数据库
                    int res = processService.saveProcess(process);
                    if (res <= 0) {
                        allSuccess = false;
                        System.err.println("Failed to save process: " + process);
                    }
                } catch (Exception e) {
                    allSuccess = false;

                }
            }

            // 所有账户处理完成后，统一进行ACK确认
            if (allSuccess) {
                channel.basicAck(deliveryTag, false);
            } else {
                // 处理失败，根据业务需求选择拒绝消息或重试
                channel.basicReject(deliveryTag, false); // false表示不重新入队
            }
        } catch (Exception e) {
            // 处理消息解析或其他严重错误

            channel.basicReject(deliveryTag, false); // 拒绝消息，不重新入队
        }
    }

    @RabbitListener(queues = "hotfix_queue")
    public void receiveHotfix(String message, @Headers Map<String, Object> headers, Channel channel) throws IOException {
        System.out.println("Received message: " + message);

        // 解密数据
        try {
            message = EncryptionUtils.decrypt(message);
            System.out.println("Decrypted message: " + message);
        } catch (Exception e) {
            System.out.println("解密失败");
        }

        // 获取消息的deliveryTag，用于后续确认
        Long deliveryTag = (Long) headers.get(AmqpHeaders.DELIVERY_TAG);


        try {
            List<Hotfix> hotfixes = JSON.parseArray(message, Hotfix.class);
            boolean allSuccess = true;

            // 删除数据库中已存在的补丁对象

            Set<String> macAddresses = new HashSet<>();
            for (Hotfix hotfix : hotfixes) {
                macAddresses.add(hotfix.getMacAddress());
            }
            // 批量删除这些macAddress对应的记录
            for (String macAddress : macAddresses) {
                hotfixService.deleteByMacAddress(macAddress);
            }


            // 处理所有补丁对象
            GenerationResult result = Ai.hotfixAi(hotfixes);
            String analyzedJson = result.getOutput().getChoices().get(0).getMessage().getContent();
            List<Hotfix> analyzedHotfixes = JSON.parseArray(analyzedJson, Hotfix.class);
            for (Hotfix hotfix : analyzedHotfixes) {
                try {
                    // 存储到数据库
                    int res = hotfixService.saveHotfix(hotfix);
                    if (res <= 0) {
                        allSuccess = false;
                        System.err.println("Failed to save hotfix: " + hotfix);
                    }
                } catch (Exception e) {
                    allSuccess = false;

                }
            }

            // 所有补丁处理完成后，统一进行ACK确认
            if (allSuccess) {
                channel.basicAck(deliveryTag, false);
            } else {
                // 处理失败，根据业务需求选择拒绝消息或重试
                channel.basicReject(deliveryTag, false); // false表示不重新入队
            }
        } catch (Exception e) {
            // 处理消息解析或其他严重错误

            channel.basicReject(deliveryTag, false); // 拒绝消息，不重新入队
        }
    }

    @RabbitListener(queues = "weak_pwd_queue")
    public void receiveWeakPwd(String message, @Headers Map<String, Object> headers, Channel channel) throws IOException {
        System.out.println("Received message: " + message);

        // 解密数据
        try {
            message = EncryptionUtils.decrypt(message);
            System.out.println("Decrypted message: " + message);
        } catch (Exception e) {
            System.out.println("解密失败");
        }

        // 获取消息的deliveryTag，用于后续确认
        Long deliveryTag = (Long) headers.get(AmqpHeaders.DELIVERY_TAG);

        try {
            List<WeakPwd> weakPwds = JSON.parseArray(message, WeakPwd.class);
            boolean allSuccess = true;

            // 删除数据库中已存在的弱密码对象

            Set<String> macAddresses = new HashSet<>();
            for (WeakPwd  weakPwd : weakPwds ) {
                macAddresses.add(weakPwd.getMacAddress());
            }

            // 批量删除这些macAddress对应的记录
            for (String macAddress : macAddresses) {
                weakPwdService.deleteByMacAddress(macAddress);
            }


            // 处理所有弱密码对象
            for (WeakPwd weakPwd  : weakPwds) {
                try {
                    // 存储到数据库
                    int res = weakPwdService.saveWeakPwd(weakPwd);
                    if (res <= 0) {
                        allSuccess = false;
                        System.err.println("Failed to save weak_pwd: " + weakPwd);
                    }
                } catch (Exception e) {
                    allSuccess = false;

                }
            }

            // 所有弱密码处理完成后，统一进行ACK确认
            if (allSuccess) {
                channel.basicAck(deliveryTag, false);
            } else {
                // 处理失败，根据业务需求选择拒绝消息或重试
                channel.basicReject(deliveryTag, false); // false表示不重新入队
            }
        } catch (Exception e) {
            // 处理消息解析或其他严重错误

            channel.basicReject(deliveryTag, false); // 拒绝消息，不重新入队
        }
    }

    @RabbitListener(queues = "app_vul_queue")
    public void receiveAppVul(String message, @Headers Map<String, Object> headers, Channel channel) throws IOException {
        System.out.println("Received message: " + message);

        // 解密数据
        try {
            message = EncryptionUtils.decrypt(message);
            System.out.println("Decrypted message: " + message);
        } catch (Exception e) {
            System.out.println("解密失败");
        }

        // 获取消息的deliveryTag，用于后续确认
        Long deliveryTag = (Long) headers.get(AmqpHeaders.DELIVERY_TAG);

        try {
            List<HostAppVul> hostAppVuls = JSON.parseArray(message, HostAppVul.class);
            boolean allSuccess = true;

            // 删除数据库中已存在的弱密码对象

            Set<String> macAddresses = new HashSet<>();
            for (HostAppVul  hostAppVul : hostAppVuls ) {
                macAddresses.add(hostAppVul.getMacAddress());
            }

            // 批量删除这些macAddress对应的记录
            for (String macAddress : macAddresses) {
                hostAppVulService.deleteByMacAddress(macAddress);
            }


            // 处理所有弱密码对象
            for (HostAppVul hostAppVul  : hostAppVuls) {
                try {
                    // 存储到数据库
                    int res = hostAppVulService.saveHostAppVul(hostAppVul);
                    if (res <= 0) {
                        allSuccess = false;
                        System.err.println("Failed to save host_app_vul: " + hostAppVul);
                    }
                } catch (Exception e) {
                    allSuccess = false;

                }
            }

            // 所有弱密码处理完成后，统一进行ACK确认
            if (allSuccess) {
                channel.basicAck(deliveryTag, false);
            } else {
                // 处理失败，根据业务需求选择拒绝消息或重试
                channel.basicReject(deliveryTag, false); // false表示不重新入队
            }
        } catch (Exception e) {
            // 处理消息解析或其他严重错误

            channel.basicReject(deliveryTag, false); // 拒绝消息，不重新入队
        }
    }

    @RabbitListener(queues = "vul_queue")
    public void receiveVul(String message, @Headers Map<String, Object> headers, Channel channel) throws IOException {
        System.out.println("Received message: " + message);

        // 解密数据
        try {
            message = EncryptionUtils.decrypt(message);
            System.out.println("Decrypted message: " + message);
        } catch (Exception e) {
            System.out.println("解密失败");
        }

        // 获取消息的deliveryTag，用于后续确认
        Long deliveryTag = (Long) headers.get(AmqpHeaders.DELIVERY_TAG);

        try {
            List<HostVul> hostVuls = JSON.parseArray(message, HostVul.class);
            boolean allSuccess = true;

            // 删除数据库中已存在的弱密码对象

            Set<String> macAddresses = new HashSet<>();
            for (HostVul  hostVul : hostVuls ) {
                macAddresses.add(hostVul.getMacAddress());
            }

            // 批量删除这些macAddress对应的记录
            for (String macAddress : macAddresses) {
                hostVulService.deleteByMacAddress(macAddress);
            }


            // 处理所有对象
            for (HostVul hostVul  : hostVuls) {
                try {
                    // 存储到数据库
                    int res = hostVulService.saveHostVul(hostVul);
                    if (res <= 0) {
                        allSuccess = false;
                        System.err.println("Failed to save host_vul: " + hostVul);
                    }
                } catch (Exception e) {
                    allSuccess = false;

                }
            }

            // 所有弱密码处理完成后，统一进行ACK确认
            if (allSuccess) {
                channel.basicAck(deliveryTag, false);
            } else {
                // 处理失败，根据业务需求选择拒绝消息或重试
                channel.basicReject(deliveryTag, false); // false表示不重新入队
            }
        } catch (Exception e) {
            // 处理消息解析或其他严重错误

            channel.basicReject(deliveryTag, false); // 拒绝消息，不重新入队
        }
    }


    @RabbitListener(queues = "system_vul_queue")
    public void receiveSysVul(String message, @Headers Map<String, Object> headers, Channel channel) throws IOException {
        System.out.println("Received message: " + message);

        // 解密数据
        try {
            message = EncryptionUtils.decrypt(message);
            System.out.println("Decrypted message: " + message);
        } catch (Exception e) {
            System.out.println("解密失败");
        }

        // 获取消息的deliveryTag，用于后续确认
        Long deliveryTag = (Long) headers.get(AmqpHeaders.DELIVERY_TAG);

        try {
            List<HostSystemVul> hostSystemVuls = JSON.parseArray(message, HostSystemVul.class);
            boolean allSuccess = true;

            // 删除数据库中已存在的弱密码对象

            Set<String> macAddresses = new HashSet<>();
            for (HostSystemVul  hostSystemVul : hostSystemVuls ) {
                macAddresses.add(hostSystemVul.getMacAddress());
            }

            // 批量删除这些macAddress对应的记录
            for (String macAddress : macAddresses) {
                hostSystemVulService.deleteByMacAddress(macAddress);
            }


            // 处理所有对象
            for (HostSystemVul hostSystemVul  : hostSystemVuls) {
                try {
                    // 存储到数据库
                    int res = hostSystemVulService.saveHostSystemVul(hostSystemVul);
                    if (res <= 0) {
                        allSuccess = false;
                        System.err.println("Failed to save host_system_vul: " + hostSystemVul);
                    }
                } catch (Exception e) {
                    allSuccess = false;

                }
            }

            // 所有弱密码处理完成后，统一进行ACK确认
            if (allSuccess) {
                channel.basicAck(deliveryTag, false);
            } else {
                // 处理失败，根据业务需求选择拒绝消息或重试
                channel.basicReject(deliveryTag, false); // false表示不重新入队
            }
        } catch (Exception e) {
            // 处理消息解析或其他严重错误

            channel.basicReject(deliveryTag, false); // 拒绝消息，不重新入队
        }
    }

    @RabbitListener(queues = "logs_queue")
    public void receiveHostLog(String message, @Headers Map<String, Object> headers, Channel channel) throws IOException {
        System.out.println("Received message: " + message);

        // 解密数据
        try {
            message = EncryptionUtils.decrypt(message);
            System.out.println("Decrypted logmessage: " + message);
        } catch (Exception e) {
            System.out.println("解密失败");
        }

        // 获取消息的deliveryTag，用于后续确认
        Long deliveryTag = (Long) headers.get(AmqpHeaders.DELIVERY_TAG);

        try {
            List<HostLog> hostLogs = JSON.parseArray(message, HostLog.class);
            boolean allSuccess = true;
            Set<String> macAddresses = new HashSet<>();
//            for (HostLog  hostLog : hostLogs ) {
//                macAddresses.add(hostLog.getMacAddress());
//            }
//
//            // 批量删除这些macAddress对应的记录
//            for (String macAddress : macAddresses) {
//                hostLogService.deleteByMacAddress(macAddress);
//            }


            // 处理所有对象
            GenerationResult result = Ai.hostLogAi(hostLogs);
            String analyzedJson = result.getOutput().getChoices().get(0).getMessage().getContent();
            List<HostLog> analyzedHostLogs = JSON.parseArray(analyzedJson, HostLog.class);
            for (HostLog hostLog : analyzedHostLogs) {
                try {
                    // 存储到数据库
                    int res = hostLogService.saveHostLog(hostLog);
                    if (res <= 0) {
                        allSuccess = false;
                        System.err.println("Failed to save hostLog: " + hostLog);
                    }
                } catch (Exception e) {
                    allSuccess = false;

                }
            }

            // 所有弱密码处理完成后，统一进行ACK确认
            if (allSuccess) {
                channel.basicAck(deliveryTag, false);
            } else {
                // 处理失败，根据业务需求选择拒绝消息或重试
                channel.basicReject(deliveryTag, false); // false表示不重新入队
            }
        } catch (Exception e) {
            // 处理消息解析或其他严重错误

            channel.basicReject(deliveryTag, false); // 拒绝消息，不重新入队
        }
    }
}
